import torch
from torchvision import transforms
from models import crnn
import matplotlib.pyplot as plt
from PIL import Image

class resizeNormalize(object):
	def __init__(self, size, interpolation=Image.BILINEAR):
		self.size = size
		self.interpolation = interpolation
		self.toTensor = transforms.ToTensor()

	def __call__(self, img):
		img = img.resize(self.size, self.interpolation)
		img = self.toTensor(img)
		img.sub_(0.5).div_(0.5)
		return img

def visualize_features(model, image):
    model.eval()
    with torch.no_grad():
        features = model.cnn(image.unsqueeze(0))
        plt.imshow(features[0, 0].cpu().numpy(), cmap='jet')
        plt.title("CNN Features")
        plt.show()

# 随机抽取一张图片测试

data = Image.open("D:/codework/python_work/ocr_work/crnn-ocr/captcha-ocr/DataSet/images/1.png").convert('L')
(w,h) = data.size
size_h = 32
ratio = 32 / float(h)
size_w = int(w * ratio)
transform = resizeNormalize((size_w,size_h))
image = transform(data)
char_set = open('chars.txt', 'r', encoding='utf-8').readlines()
char_set = ''.join([ch.strip('\n') for ch in char_set[1:]] + ['卍'])
n_class = len(char_set)
torch.set_printoptions(threshold=float('inf'))
model = crnn.CRNN(32, 1, n_class, 256)
modelpath = "pytorch-crnn.pth"
if os.path.exists(modelpath):
    print('Load model from "%s" ...' % modelpath)
    model.load_state_dict(torch.load(modelpath, map_location=torch.device("cpu")))
    print('Done!')
model.eval()
visualize_features(model, image)